#include <cstdio>
#include <vector>
#include <cuda_runtime.h>

__device__ __forceinline__ float2 ComplexMul(float2 a, float2 b) {
  return make_float2(a.x * b.x - a.y * b.y, a.x * b.y + a.y * b.x);
}
__device__ __forceinline__ float2 ComplexSq(float2 a) {
  return make_float2(a.x * a.x - a.y * a.y, 2 * a.x * a.y);
}
__device__ __forceinline__ float Dot(float2 a, float2 b) {
  return a.x*b.x + a.y*b.y;
}
__device__ float ComputeDistance(float2 c, int bailout) {
  float2 z = make_float2(0, 0);
  float2 dz = make_float2(1, 0);
  for (int i = 0; i < bailout; ++i) {
    dz = ComplexMul(z, dz);
    dz.x = 2.0f * dz.x + 1.0f;
    dz.y = 2.0f * dz.y;
    z = ComplexSq(z);
    z.x += c.x;
    z.y += c.y;
    const float m2 = Dot(z, z);
    if (m2 > 100.0f) {
      return sqrtf(m2 / Dot(dz, dz)) * 0.5f * logf(m2);
    }
  }
  return 0.0f;
}
__global__ void Mandelbrot(float2 center, float zoom, int bailout, int size, uchar4 *image) {
  const int x = blockIdx.x * blockDim.x + threadIdx.x;
  const int y = blockIdx.y * blockDim.y + threadIdx.y;
  if (x < size && y < size) {
    const int idx = x + y * size;

    float2 c = make_float2(2.0f * (x / static_cast<float>(size) - 0.5f),
                           2.0f * (y / static_cast<float>(size) - 0.5f));
    c.x = c.x / zoom + center.x;
    c.y = c.y / zoom + center.y;
    float d = 255.0f * sqrtf(sqrtf(ComputeDistance(c, bailout) * zoom));

    image[idx] = make_uchar4(d, d, d, 255);
  }
}

int main(int argc, char *argv[]) {
  float2 center = make_float2(-0.761574f, -0.0847596f);
  float zoom = 1000.0f;
  int bailout = 1000;
  int size = 800;

  for (int i = 0; i < argc; ++i) {
    if (strcmp(argv[i], "-re") == 0) {
      center.x = atof(argv[++i]);
    } else if (strcmp(argv[i], "-im") == 0) {
      center.y = atof(argv[++i]);
    } else if (strcmp(argv[i], "-zoom") == 0) {
      zoom = atof(argv[++i]);
    } else if (strcmp(argv[i], "-bailout") == 0) {
      bailout = atoi(argv[++i]);
    } else if (strcmp(argv[i], "-size") == 0) {
      size = atoi(argv[++i]);
    }
  }
  printf("re      = %f\n", center.x);
  printf("im      = %f\n", center.y);
  printf("zoom    = %f\n", zoom);
  printf("bailout = %d\n", bailout);
  printf("size    = %d\n", size);

  uchar4 *image_dev;
  cudaMalloc(&image_dev, size * size * sizeof(uchar4));

  dim3 block(16, 16);
  dim3 grid((size + block.x - 1) / block.x, (size + block.y - 1) / block.y);
  Mandelbrot<<<grid, block>>>(center, zoom, bailout, size, image_dev);

  std::vector<uchar4> image_host(size * size);
  cudaMemcpy(&image_host[0], image_dev,
             size * size * sizeof(uchar4),
             cudaMemcpyDeviceToHost);

  FILE *fp = fopen("mandelbrot.tga", "wb");
  const uchar1 tga_header[18] = {
    0, 0, 2, 0, 0, 0, 0, 0, 0, 0, 0, 0,
    size & 0x00ff, (size & 0xff00) >> 8,
    size & 0x00ff, (size & 0xff00) >> 8,
    32, 0
  };
  fwrite(tga_header, sizeof(tga_header), 1, fp);
  fwrite(&image_host[0], sizeof(uchar4) * image_host.size(), 1, fp);
  fclose(fp);

  return 0;
}
